import pandas as pd
from tqdm import tqdm
import random


def process():
    data = pd.read_csv('data/train.txt', names=['dataid', 'classes', 'ori', 'entailment'], sep='\t')

    def smp(sam):
        classes = sam.classes
        dataid = sam.dataid
        comd = data.loc[data.classes == classes]
        negd = data.loc[data.classes != classes]
        sample = comd.sample(n=1)
        sdataid = sample.dataid.values[0]
        if dataid != sdataid:
            return sample.ori.values[0], negd.sample(n=1).ori.values[0]

    data[['contradiction1', 'contradiction2']] = data.apply(smp, axis=1, result_type='expand')
    data.to_csv('data/train.csv', index=None, header=None, sep='\t')


def process1():
    data = pd.read_csv('resd.txt', names=['dataid', 'classes', 'ori', 'entailment'], sep='\t')

    def smp(sam):
        classes = sam.classes
        dataid = sam.dataid
        comd = data.loc[data.classes == classes]
        negd = data.loc[data.classes != classes]
        sample = comd.sample(n=1)
        sdataid = sample.dataid.values[0]
        if dataid != sdataid:
            return sample.ori.values[0], negd.sample(n=1).ori.values[0]
        else:
            sample = comd.sample(n=1)
            return sample.ori.values[0], negd.sample(n=1).ori.values[0]

    with open('train.txt', 'a+', encoding='utf-8') as f:
        for idx, row in tqdm(data.iterrows()):
            f.write('\t'.join([row.dataid, row.classes, row.ori, row.entailment, *smp(row)]) + '\n')


def process2():
    res = open('addtrain.txt', 'a+', encoding='utf-8')
    with open('../data/alltraindata11.txt', 'r', encoding='utf-8') as f:
        alldata = f.readlines()
        alllen = len(alldata)
        for idx, i in tqdm(enumerate(alldata)):
            line = i.strip().split('\t')
            line = line[:4]
            if len(line) != 4: continue
            if len(line[1]) != 2: continue
            classes = line[1]
            ocls = classes
            while classes == ocls:
                oidx = random.randint(10000, alllen - 10000)
                oline = alldata[oidx].strip().split('\t')
                if len(oline) != 4: continue
                if len(oline[1]) != 2: continue
                ocls = oline[1]
            oneg = oline[2]
            line.append(oneg)
            res.write('\t'.join(line) + '\n')


def splitdata(path='../data/simtitle.txt', ratio=0.02):
    data = pd.read_csv(path,
                       names=['dataid', 'classes', 'ori', 'entailment', 'contradiction'],
                       sep='\t')
    classes = data.classes.unique()
    devdata = pd.DataFrame(columns=['dataid', 'classes', 'ori', 'entailment', 'contradiction'])
    for cls in classes:
        clsdata = data.loc[data.classes == cls]
        dev = clsdata.sample(frac=ratio)
        devdata = pd.concat([devdata, dev])

    traindata = data.loc[~data.index.isin(devdata.index)]
    traindata = traindata.sample(frac=1)
    devdata = devdata.sample(frac=1)
    traindata.to_csv('../data/traindata.csv', index=None, header=None, sep='\t')
    devdata.to_csv('../data/devdata.csv', index=None, header=None, sep='\t')


# def process3():
#     with open('data/train.txt','r',encoding='utf-8') as f:
#         for line in f:

if __name__ == '__main__':
    splitdata()
# data = pd.read_csv('data/devdata.csv', names=['dataid', 'classes', 'ori', 'entailment', 'contradiction1', 'contradiction2'],
#                    sep='\t')
# data = data.loc[:, ['ori', 'entailment', 'contradiction1', 'contradiction2']].values
# for i in data:
#     print(len(i))
